import h5py
from typing import Dict, Optional, Any

class DataKeys:
    """HDF5文件中的数据键映射"""
    MAPPINGS = {
        "DTL-010:RFS-DIG-101:Dwn0-XAxis": "axis",
        "MEBT-010:RFS-DIG-101:Dwn0-Cmp0": "mag_cav",
        "MEBT-010:RFS-DIG-101:Dwn0-Cmp1": "phase_cav",
        "MEBT-010:RFS-DIG-101:Dwn5-Cmp0": "mag_for",
        "MEBT-010:RFS-DIG-101:Dwn5-Cmp1": "phase_for",
        "MEBT-010:RFS-DIG-101:Dwn6-Cmp0": "mag_refl",
        "MEBT-010:RFS-DIG-101:Dwn6-Cmp1": "phase_refl"
    }

def extract_dataset(group: h5py.Group, index: int) -> Optional[Any]:
    """从HDF5组中提取数据集
    
    Args:
        group: HDF5数据组
        index: 时间戳索引
        
    Returns:
        Optional[Any]: 提取的数据集，如果不存在则返回None
    """
    if len(group) == 0:
        return None
        
    if index >= len(group):
        print(f"警告：索引 {index} 超出范围，组中只有 {len(group)} 个数据集")
        return None
        
    time_stamps = list(group)
    return group[time_stamps[index]][()]

def read_hdf5_find_powers(hdf5_filename: str, index: int) -> Dict[str, Any]:
    """从HDF5文件中读取功率数据
    
    Args:
        hdf5_filename: HDF5文件路径
        index: 时间戳索引
        
    Returns:
        Dict[str, Any]: 包含功率数据的字典
        
    Raises:
        FileNotFoundError: 当文件不存在时
        ValueError: 当文件格式不正确时
    """
    try:
        raw_powers = {}
        with h5py.File(hdf5_filename, 'r') as hdf5_file:
            for h5_key, raw_key in DataKeys.MAPPINGS.items():
                if h5_key not in hdf5_file:
                    print(f"警告：文件中缺少键 {h5_key}")
                    continue
                    
                group = hdf5_file[h5_key]
                if raw_key == "axis":
                    raw_powers[raw_key] = extract_dataset(group, 0)
                else:
                    raw_powers[raw_key] = extract_dataset(group, index)
                    
        return raw_powers
        
    except FileNotFoundError:
        raise FileNotFoundError(f"找不到文件：{hdf5_filename}")
    except Exception as e:
        raise ValueError(f"读取文件时出错：{str(e)}")

if __name__ == '__main__':
    # 测试代码
    try:
        test_file = "path/to/test.hdf5"
        data = read_hdf5_find_powers(test_file, 1)
        print("成功读取数据：", list(data.keys()))
    except Exception as e:
        print(f"测试失败：{str(e)}")        